
#include "ephi.hpp"

#define PATH_SUFIX ".png"

class PlasmaOctreeBuilder : public OctreeBuilder
{
public:
  PlasmaOctreeBuilder (const vect3d& c0, const vect3d& c1, prec_t resolution)
    : OctreeBuilder (c0, c1, 1, 2, 100, resolution), pmap(0)
  {
  }

  void reset (PMap& pmap) { this->pmap = &pmap; }

  void addValues (const vect3d& pos, CubicScalar *scalars, CubicVector *vectors)
  {
    pmap->addCubicValues (pos, scalars[0], vectors[0], vectors[1]);
  }

  bool inContact (const vect3d& pos)
  {
    return false;
  }

private:
  PMap *pmap;
};


Plasma::Plasma (etree_t tree_type) : dyn(statics), pfield (PField::POINTSYSTEM), scene(dyn), tree_type(tree_type)
{
  qmul = 1;
  pmap_box = 0.5;
  pmap_resolution = 0.05;
  pmap_smooth = 0.15;
  stepmag = 1;
  octree_resolution = 0.05;
  octree_c0 = vect3d(-pmap_box, -pmap_box, -pmap_box);
  octree_c1 = vect3d(pmap_box, pmap_box, pmap_box);
  particle_plane = Screen::PLANE_XY;
  dyn.setDisableFields (true);
  calibrate_type = NONE;
  octree_radius = 0;
  builders[0] = 0;
  bmul = emul = 1.0;
}

Plasma::~Plasma ()
{
  for ( size_t i = 0; i < 3; i++ )
      delete builders[i];
}

void Plasma::init ()
{
  if ( !builders[0] )
    {
      dyntree = new Dynamic3 ();
      switch (tree_type)
        {
          case TREE_FULL: {
            dyntree->get(0).addStaticElement (new StaticOctree (trees[0], true, true, true, bmul, emul));
            dyntree->get(1).addStaticElement (new StaticOctree (trees[1], true, true, true, bmul, emul));
            dyntree->get(2).addStaticElement (new StaticOctree (trees[2], true, true, true, bmul, emul));
            break;
          }
          case TREE_24: {
            dyntree->get(0).addStaticElement (new StaticOctree24 (trees[0], true, true, true, bmul, emul));
            dyntree->get(1).addStaticElement (new StaticOctree24 (trees[1], true, true, true, bmul, emul));
            dyntree->get(2).addStaticElement (new StaticOctree24 (trees[2], true, true, true, bmul, emul));

            planes.clear();
            planes.push_back (vect3d(0, 1, 0));
            planes.push_back (vect3d(0, 0, 1));
            planes.push_back (vect3d (1, -1, 0).normal());
            planes.push_back (vect3d (1, 0, -1).normal());

            std::vector<transf3d> transforms;
            transforms.reserve (24);
            const transf3d_pair *tprs = get_cube_polywell_transforms24 ();
            for ( size_t i = 0; i < 24; i++ )
              {
                symm.push_back (tprs[i].f);
                transforms.push_back (tprs[i].f);
              }
            pfield.setTransforms (planes, transforms);
            break;
          }
          case TREE_48: {
            dyntree->get(0).addStaticElement (new StaticOctree48 (trees[0], true, true, true, bmul, emul));
            dyntree->get(1).addStaticElement (new StaticOctree48 (trees[1], true, true, true, bmul, emul));
            dyntree->get(2).addStaticElement (new StaticOctree48 (trees[2], true, true, true, bmul, emul));

            planes.clear();
            planes.push_back (vect3d(0, 1, 0));
            planes.push_back (vect3d(1, 0, -1).normal());
            planes.push_back (vect3d(0, -1, 1).normal());

            std::vector<transf3d> transforms;
            transforms.reserve (48);
            const transf3d_pair *tprs = get_cube_polywell_transforms48 ();
            for ( size_t i = 0; i < 48; i++ )
              {
                if ( !(i & 1) )
                    rsymm.push_back (tprs[i].f);
                else
                    symm.push_back (tprs[i].f);
                transforms.push_back (tprs[i].f);
              }
            pfield.setTransforms (planes, transforms);

            break;
          }
        }
      dyn.addDynamicElement (dyntree);

      for ( size_t j = 0; j < 3; j++ )
        {
          builders[j] = new PlasmaOctreeBuilder (octree_c0, octree_c1, octree_resolution);
          builders[j]->setLimitRadius (octree_radius);
          for ( size_t i = 0; i < planes.size(); i++ )
              builders[j]->addLimit (planes[i]);
        }
    }

  fieldupdate();
  dyntree->next();
  fieldupdate();
  dyntree->next();
  fieldupdate();
  dyn.reset();
}

void Plasma::step ()
{
  timer t;
  for ( size_t i = 0; i < stepmag; i++ )
      dyn.step();
  t.report ("step");
  dyntree->next();
  fieldupdate();
  dyn.reset();
}

void Plasma::fieldupdate ()
{
  PMap pmap(vect3d(pmap_box, pmap_box, pmap_box), pmap_resolution, 0, pmap_smooth);
  fill (pmap);
  if ( !pfield.empty() )
    {
      // we use a very crude approximation. it's probably good enough for the purpouse though
      PMap potmap(vect3d(pmap_box, pmap_box, pmap_box), pmap_box / 10.0);
      fill (potmap);
      potmap.recalc();
      Statics st;
      st.addStaticElement (new StaticPMapCrudeEfield (potmap));
      pfield.setStatics (st);
      pfield.equalize ();
      pfield.fill (pmap);
    }
  pmap.recalc();
  builders[dyntree->head()]->reset (pmap);
  builders[dyntree->head()]->rebuild (trees[dyntree->head()]);
}

void Plasma::fill (PMap& pmap)
{
  if ( qmul == 0 )
      return;
  //printf ("Fill works!\n");

  if ( symm.empty() )
    {
      for ( size_t i = 0; i < dyn.size(); i++ )
        {
          if ( !pmap.addQ (dyn.get_pos(i), dyn.get_v(i), dyn.get_q(i) * qmul) )
            {
              char b[128];
              Ephi::debug (Ephi::DEBUG_APP, "Erasing particle %u at %s\n", (unsigned int)i, dyn.get_pos(i).sprint(b));
              dyn.erase (i);
              i--;
            }
        }
    }
  else
    {
      for ( size_t i = 0; i < dyn.size(); i++ )
        {
          for ( size_t j = 0; j < symm.size(); j++ )
              pmap.addQ (symm[j] * dyn.get_pos(i), symm[j] * dyn.get_v(i), dyn.get_q(i) * qmul);
          for ( size_t j = 0; j < rsymm.size(); j++ )
              pmap.addQ (rsymm[j] * dyn.get_pos(i), -(rsymm[j] * dyn.get_v(i)), dyn.get_q(i) * qmul);
        }
    }
}

void Plasma::render (const char *name, bool do_cross)
{
  std::string prefix = name;
  if ( !pfield.empty() )
    {
      std::string str = pfield.getChargeConcentrationScript ();
      FILE *fp = fopen ((prefix + "_charge.pov").c_str(), "w");
      fprintf (fp, "camera { location <3, 3, 2>*0.19 look_at <0, 0, 0> sky <0, 0, 1> right x*image_width/image_height }\n");
      fprintf (fp, "#include \"scene.inc\"\n");
      fprintf (fp, "%s", str.c_str());
      fclose (fp);

      str = pfield.getPovrayScript ();
      fp = fopen ((prefix + "_polywell.pov").c_str(), "w");
      fprintf (fp, "camera { location <3, 3, 2>*0.19 look_at <0, 0, 0> sky <0, 0, 1> right x*image_width/image_height }\n");
      fprintf (fp, "#include \"scene.inc\"\n");
      fprintf (fp, "%s", str.c_str());
      fclose (fp);
    }

  Screen screen(800, 800, 2 * pmap_box, 2 * pmap_box);
  screen.set_axis (particle_plane);
  render_screen (scene, screen, prefix, "_");
  if ( do_cross )
    {
      screen.set_axis (vect3d(0, 0, 0), vect3d(1, 1, 0).normal(), vect3d(0, 0, 1));
      render_screen (scene, screen, prefix, "_X");
    }
}

void Plasma::render_screen (Scene& scene, Screen& screen, const std::string& name, const std::string& prefix)
{
  scene.calc (screen, Scene::CALC_EFIELD);
  scene.set_coloring (0, false);
  calibrate (scene, prefix + "efield" PATH_SUFIX);
  scene.render_map (screen);
  screen.write (name + prefix + "efield" PATH_SUFIX);
  scene.set_coloring (32, true);
  calibrate (scene, prefix + "efield_log_stepped" PATH_SUFIX);
  scene.render_map (screen);
  screen.write (name + prefix + "efield_log_stepped" PATH_SUFIX);
  scene.set_coloring (0, true);
  calibrate (scene, prefix + "efield_lic" PATH_SUFIX);
  scene.render_lic (screen);
  screen.write (name + prefix + "efield_lic" PATH_SUFIX);
  scene.calc (screen, Scene::CALC_BFIELD);
  scene.set_coloring (0, false);
  calibrate (scene, prefix + "bfield" PATH_SUFIX);
  scene.render_map (screen);
  screen.write (name + prefix + "bfield" PATH_SUFIX);
  scene.set_coloring (32, true);
  calibrate (scene, prefix + "bfield_log_stepped" PATH_SUFIX);
  scene.render_map (screen);
  screen.write (name + prefix + "bfield_log_stepped" PATH_SUFIX);
  scene.set_coloring (0, true);
  calibrate (scene, prefix + "bfield_lic" PATH_SUFIX);
  scene.render_lic (screen);
  screen.write (name + prefix + "bfield_lic" PATH_SUFIX);
  scene.calc (screen, Scene::CALC_POTENTIAL);
  scene.set_coloring (32, false);
  calibrate (scene, prefix + "pot" PATH_SUFIX);
  scene.render_map (screen);
  screen.write (name + prefix + "pot" PATH_SUFIX);

  screen.clear();
  statics.draw (screen);
  screen.draw_axes();
  prec_t trail_len = screen.get_space_width() / 100;
  for ( size_t i = 0; i < dyn.size(); i++ )
    {
      if ( symm.empty() )
        {
          vect3d pos = dyn.get_pos(i);
          vect3d v = dyn.get_v(i);
          if ( v.magnitude() == 0 || 1 )
              screen.set_pixel (pos, Screen::RED);
          else
              screen.draw_line (pos, pos + v.normal() * trail_len, Screen::RED);
        }
      else
        {
          for ( size_t j = 0; j < symm.size(); j++ )
            {
              vect3d pos = symm[j] * dyn.get_pos(i);
              vect3d v = symm[j] * dyn.get_v(i);
              if ( v.magnitude() == 0 || 1 )
                  screen.set_pixel (pos, Screen::RED);
              else
                  screen.draw_line (pos, pos + v.normal() * trail_len, Screen::RED);
            }

          for ( size_t j = 0; j < rsymm.size(); j++ )
            {
              vect3d pos = symm[j] * dyn.get_pos(i);
              vect3d v = symm[j] * dyn.get_v(i);
              if ( v.magnitude() == 0  || 1)
                  screen.set_pixel (pos, Screen::RED);
              else
                  screen.draw_line (pos, pos - v.normal() * trail_len, Screen::RED);
            }
        }
    }
#if 0
  screen.draw_line (vect3d(0, 0, 0), vect3d(octree_resolution, 0, 0), Screen::GREEN);
  screen.draw_line (vect3d(0, 0, 0), vect3d(0, octree_resolution, 0), Screen::GREEN);
  screen.draw_line (vect3d(0, 0.001, 0), vect3d(2 * pmap_smooth, 0.001, 0), Screen::BLUE);
#endif
  screen.write (name + prefix + "particles" PATH_SUFIX);
}

void Plasma::calibrate (Scene& scene, const std::string& name)
{
  if ( calibrate_type == INITIAL )
    {
      calibration_t::iterator fit = calibration.find (name);
      if ( fit == calibration.end() )
        {
          calibration.insert (calibration_t::value_type (name,
              calibr(scene.get_fmin(), scene.get_fmax(), scene.get_fadd())));
        }
      else
        {
          scene.calibrate_field (fit->second.fmin, fit->second.fmax, fit->second.fadd);
        }
    }
}

